import torch
import torch.nn as nn
import sys

sys.path.append("./")
from model.semanticview import SemanticView
from model.radicalview import RadicalViewFeature
from model.gazview import GazViewFeature
from model.synonymview import SynoymViewFeature
from loguru import logger
from utils.data import Data
from model.config_layer import config_encoder


class MultiViewEncoder(nn.Module):
    def __init__(self, data: Data):
        super(MultiViewEncoder, self).__init__()

        logger.info("Build MultiViewEncoder...")

        self.gpu = True
        self.use_radical = config_encoder["use_radical"]
        self.use_lexicon = config_encoder["use_lexicon"]
        self.use_synoym = config_encoder["use_synoym"]

        self.multiview_feature_dim = config_encoder["multiview_feature_dim"]  # TODO
        self.hidden_dim = config_encoder["hidden_dim"]
        self.pooler_dim = config_encoder["pooler_dim"]

        self.semantic_view = SemanticView(data)  # [b, l , 768 + 50*2]

        if self.use_lexicon:
            logger.info("   +LexiconView Feature..")
            self.lexicon_view = GazViewFeature(data)
            self.multiview_feature_dim = self.multiview_feature_dim + self.pooler_dim

        if self.use_radical:
            logger.info("   +RadicalView Feature..")
            self.radical_view = RadicalViewFeature(data)
            self.multiview_feature_dim = self.multiview_feature_dim + self.pooler_dim

        if self.use_synoym:
            logger.info("   +SynoymView Feature..")
            self.synoym_view = SynoymViewFeature(data)
            self.multiview_feature_dim = self.multiview_feature_dim + self.pooler_dim

        data.multiview_feature_dim = self.multiview_feature_dim

    def forward(
        self,
        batch_word,
        batch_biword,
        layer_gaz,
        gaz_count,
        gaz_chars,
        gaz_mask,
        gazchar_mask,
        batch_bert,
        bert_mask,
    ):

        multiview_feature = self.semantic_view(
            batch_word, batch_biword, batch_bert, bert_mask
        )

        if self.use_lexicon:
            lexicon_feature = self.lexicon_view(
                batch_word,
                batch_biword,
                layer_gaz,
                gaz_count,
                gaz_chars,
                gaz_mask,
                gazchar_mask,
            )
            multiview_feature = torch.cat([multiview_feature, lexicon_feature], dim=-1)

        if self.use_radical:
            radical_feature = self.radical_view(batch_word)
            multiview_feature = torch.cat([multiview_feature, radical_feature], dim=-1)

        if self.use_synoym:
            # TODO!
            # synoym_feature = self.synoym_view()
            # multiview_feature = torch.cat([multiview_feature, synoym_feature], dim=-1)
            pass

        return multiview_feature
